#!/usr/bin/env python
#################################################################################
# Copyright 2018 ROBOTIS CO., LTD.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################

# Authors: Gilbert #

import rospy
import os
import json
import numpy as np
import random
import time
import sys
sys.path.append(os.path.dirname(os.path.abspath(os.path.dirname(__file__))))
from collections import deque
from std_msgs.msg import Float32MultiArray
from src.turtlebot3_dqn.simulation_environment_real import Env
from keras.models import Sequential, load_model
from keras.optimizers import RMSprop
from keras.layers import Dense, Dropout, Activation
from geometry_msgs.msg import Pose, Twist

EPISODES = 10000

class ReinforceAgent():
    def __init__(self, state_size, action_size):
        # self.pub_result = rospy.Publisher('result', Float32MultiArray, queue_size=5)
        # self.result = Float32MultiArray()

        self.load_model = True
        self.load_episode = 4500
        self.epsilon = 0.05
        self.state_size = state_size
        self.action_size = action_size
        self.episode_step = 6000
        self.target_update = 2000
        self.discount_factor = 0.9
        self.learning_rate = 0.001
        self.epsilon = 1.0
        self.epsilon_decay = 0.995
        self.epsilon_min = 0.15
        self.batch_size = 64
        self.train_start = 64
        self.memory = deque(maxlen=1000000)
        self.start_time = time.time()
        self.model = self.buildModel()
        self.target_model = self.buildModel()

        # self.updateTargetModel()

        # self.turtlebot = rospy.Publisher('gazebo/set_model_state', ModelState, queue_size=1)

        self.dirPath = os.path.dirname(os.path.realpath(__file__))
        # self.dirPath = self.dirPath.replace('turtlebot3_dqn/nodes', 'turtlebot3_dqn/save_model/stage_1_epsilon_min_0p15')
        self.dirPath = self.dirPath.replace('turtlebot3_dqn/nodes', 'turtlebot3_dqn/bestnet/stage_1_RMSprop_Layer_2_6_')
        # self.dirPath = self.dirPath.replace('turtlebot3_dqn/nodes', 'turtlebot3_dqn/bestnet/doubledqn_target_update_2_world_1_')


        if self.load_model:
            self.model.set_weights(load_model(self.dirPath+str(self.load_episode)+".h5").get_weights())
            with open(self.dirPath+str(self.load_episode)+'.json') as outfile:
                param = json.load(outfile)
                self.epsilon = param.get('epsilon')

    def buildModel(self):
        model = Sequential()
        model.add(Dense(64, input_shape=(self.state_size,), activation='relu', kernel_initializer='lecun_uniform'))
        model.add(Dense(64, activation='relu', kernel_initializer='lecun_uniform'))
        model.add(Dense(self.action_size, kernel_initializer='lecun_uniform'))
        model.add(Activation('linear'))
        model.compile(loss='mse', optimizer=RMSprop(lr=self.learning_rate, rho=0.9, epsilon=1e-06))
        model.summary()
        return model

    # def getQvalue(self, reward, next_target, done):
    #     if done:
    #         return reward
    #     else:
    #         return reward + self.discount_factor * np.amax(next_target)

    # def updateTargetModel(self):
    #     self.target_model.set_weights(self.model.get_weights())

    def getAction(self, state):
        if state[362] > 0.3:
            if np.random.rand() <= self.epsilon:
                self.q_value = np.zeros(self.action_size)
                return random.randrange(self.action_size)
            else:
                q_value = self.model.predict(state.reshape(1, len(state)))
                self.q_value = q_value
        else:
            q_value = self.model.predict(state.reshape(1, len(state)))
            self.q_value = q_value
        # print "q_value[0]", q_value[0]
        return np.argmax(q_value[0])

    # def appendMemory(self, state, action, reward, next_state, done):
    #     self.memory.append((state, action, reward, next_state, done))

    # def trainModel(self, target=False):
    #     mini_batch = random.sample(self.memory, self.batch_size)
    #     X_batch = np.empty((0, self.state_size), dtype=np.float64)
    #     Y_batch = np.empty((0, self.action_size), dtype=np.float64)
    #
    #     for i in range(self.batch_size):
    #         states = mini_batch[i][0]
    #         actions = mini_batch[i][1]
    #         rewards = mini_batch[i][2]
    #         next_states = mini_batch[i][3]
    #         dones = mini_batch[i][4]
    #
    #         q_value = self.model.predict(states.reshape(1, len(states)))
    #         self.q_value = q_value
    #
    #         if target:
    #             next_target = self.target_model.predict(next_states.reshape(1, len(next_states)))
    #
    #         else:
    #             next_target = self.model.predict(next_states.reshape(1, len(next_states)))
    #
    #         next_q_value = self.getQvalue(rewards, next_target, dones)
    #
    #         X_batch = np.append(X_batch, np.array([states.copy()]), axis=0)
    #         Y_sample = q_value.copy()
    #
    #         Y_sample[0][actions] = next_q_value
    #         Y_batch = np.append(Y_batch, np.array([Y_sample[0]]), axis=0)
    #
    #         if dones:
    #             X_batch = np.append(X_batch, np.array([next_states.copy()]), axis=0)
    #             Y_batch = np.append(Y_batch, np.array([[rewards] * self.action_size]), axis=0)
    #
    #     self.model.fit(X_batch, Y_batch, batch_size=self.batch_size, epochs=1, verbose=0)

if __name__ == '__main__':
    rospy.init_node('turtlebot3_dqn_stage_1')
    # pub_result = rospy.Publisher('result', Float32MultiArray, queue_size=5)
    # pub_get_action = rospy.Publisher('get_action', Float32MultiArray, queue_size=5)
    # result = Float32MultiArray()
    # get_action = Float32MultiArray()
    state_size = 364
    action_size = 5

    env = Env(action_size)

    agent = ReinforceAgent(state_size, action_size)
    global_step = 0

    for e in range(agent.load_episode + 1, EPISODES):
        done = False
        state = env.reset()
        score = 0
        for t in range(agent.episode_step):
            # t = True
            # v = Twist()
            # v.linear.x = 0
            # while t:
            #     pub_cmd_vel.publish(v)
            action = agent.getAction(state)
            next_state, reward, done = env.step(action)
            score += reward
            if done:
                print "crash"
                break
            state = next_state
            global_step += 1
